#ifndef HCCLV1_GEN_TASK_NODE_UTILS_H
#define HCCLV1_GEN_TASK_NODE_UTILS_H

#include <map>
#include <set>
#include <vector>
#include <memory>

#include "log.h"
#include "task_graph_generator.h"
#include "llt_common.h"
namespace hccl {

class GenTaskNodeGraphBase {
public:
    virtual ~GenTaskNodeGraphBase();
    virtual TaskNodePtr GenGraph();

protected:
    TaskNodePtr head = nullptr;

    void LinkNode(TaskNodePtr parent, TaskNodePtr node);
    TaskNodePtr CreateHeadNode();

    TaskNodePtr CreatePostNode(RankId myRank, RankId remoteRankId, LinkInfo &link, u32 topicId,
        u32 queueId, u32 pos, Stream& stream, NotifyTypeStub notifyType = NotifyTypeStub::INVALID);
    TaskNodePtr CreateWaitNode(RankId myRank, RankId remoteRankId, LinkInfo &link, u32 topicId,
        u32 queueId, u32 pos, Stream& stream, NotifyTypeStub notifyType = NotifyTypeStub::INVALID);
    TaskNodePtr CreateLocalPostToNode(u32 topicId, RankId myRank, u32 queueId, u32 pos, Stream& stream);
    TaskNodePtr CreateLocalWaitFromNode(u32 topicId, RankId myRank, u32 queueId, u32 pos, Stream& stream);
    TaskNodePtr CreateLocalCopyNode(DataSlice& srcSlice, DataSlice& dstSlice, RankId rankId, u32 queueId, u32 pos, Stream& stream);
    TaskNodePtr CreateLocalReduceNode(DataSlice& srcSlice, DataSlice& dstSlice, RankId rankId,
        u32 queueId, u32 pos, HcclDataType dataType, HcclReduceOp op, Stream& stream);
    TaskNodePtr CreateReadNode(DataSlice& localSlice, DataSlice& remoteSlice, RankId remoteRankId,
        LinkInfo& link, RankId rankId, u32 queueId, u32 pos, Stream& stream);
    TaskNodePtr CreateWriteNode(DataSlice& localSlice, DataSlice& remoteSlice, RankId remoteRankId,
        LinkInfo& link, RankId rankId, u32 queueId, u32 pos, Stream& stream);
    TaskNodePtr CreateReadReduceNode(DataSlice& localSlice, DataSlice& remoteSlice, RankId remoteRankId,
        LinkInfo& link, HcclDataType dataType, HcclReduceOp op, RankId rankId, u32 queueId, u32 pos, Stream& stream);
    TaskNodePtr CreateWriteReduceNode(DataSlice& localSlice, DataSlice& remoteSlice, RankId remoteRankId,
        LinkInfo& link, HcclDataType dataType, HcclReduceOp op, RankId rankId, u32 queueId, u32 pos, Stream& stream);
private:
    std::vector<TaskNodePtr> toDeleteTaskNodeResource_;
};

}

#endif